`include "cpu_def.vh"

module lsu(
  input clk,
  input rst,

  // from ex stage
  input                   lsu_valid_i    ,
  input                   lsu_read_mem   ,
  input                   lsu_write_mem  ,
  input [`NR_LD_OP - 1:0] lsu_ld_op      ,
  input [`NR_ST_OP - 1:0] lsu_st_op      ,
  input [           31:0] lsu_addr       ,
  input [           31:0] lsu_rt         ,
  input                   lsu_ex_taken   ,

  output lsu_stall_ex,

  // to mmu
  output [31:0] data_vaddr,
  output        data_vaddr_valid,
  output        is_store  ,

  // from mmu
  input [31:0] data_paddr       ,
  input        data_paddr_valid ,
  input        data_uncache     ,
  input        data_tlb_refill  ,
  input        data_tlb_invalid ,
  input        data_tlb_modified,

  // data sram interface
  (*mark_debug = "true"*)output data_sram_req    ,
  (*mark_debug = "true"*)input  data_sram_addr_ok,

  (*mark_debug = "true"*)output        data_sram_uncache,
  (*mark_debug = "true"*)output        data_sram_wr     ,
  (*mark_debug = "true"*)output [ 1:0] data_sram_size   ,
  (*mark_debug = "true"*)output [31:0] data_sram_addr   ,
  (*mark_debug = "true"*)output [ 3:0] data_sram_wstrb  ,
  (*mark_debug = "true"*)output [31:0] data_sram_wdata  ,

  (*mark_debug = "true"*)input        data_sram_data_ok,
  (*mark_debug = "true"*)input [31:0] data_sram_rdata  ,

  // to wb stage
  output        lsu_valid_o,
  output [31:0] lsu_ld_data,

  output lsu_stall_wb,

  // to ex stage
  output        lsu_adel    ,
  output        lsu_ades    ,
  output        lsu_tlbrel  ,
  output        lsu_tlbres  ,
  output        lsu_tlbil   ,
  output        lsu_tlbis   ,
  output        lsu_tlbm    ,
  output [31:0] lsu_badvaddr
);

  //                  +--+
  //                  |  |
  //                  |WB|
  //                  |/\|
  //                  +--+
  // +--+             +--+
  // |  |  +-------+  |  |  +-------+
  // |EX|->|stage 0|->|R1|->|stage 1|
  // |/\|  +-------+  |/\|  +-------+
  // +--+      |      +--+      ^
  //           |      +--+      |
  //           ------>|d$|-------
  //                  +--+

  /* stage 0: request */
  reg data_sram_addr_ok_r;

  wire [31:0] addr   = lsu_addr;
  wire [31:0] rdata  = data_sram_rdata;
  wire [ 1:0] offset = addr[1:0];
  wire [31:0] lsu_st_result [`NR_ST_OP - 1:0];

  wire                        s0_out_valid;
  wire [`NR_LD_OP -      1:0] s0_out_ld_op;
  wire [                 1:0] s0_out_offset;
  wire [                31:0] s0_out_rt;
  wire [34 + `NR_LD_OP - 1:0] s0_out;

  /* stage 1: response */
  reg                        s1_valid;
  reg [34 + `NR_LD_OP - 1:0] s1_reg;

  wire [`NR_LD_OP - 1:0] s1_ld_op;
  wire [            1:0] s1_offset;
  wire [           31:0] s1_rt;
  wire                   s1_stall;
  wire                   s1_flush;

  wire [31:0] lsu_ld_result [`NR_LD_OP - 1:0];

  /* stage 0 */
  assign lsu_stall_ex  = 
    (data_vaddr_valid && !data_paddr_valid) || // tlb lookup 
    (data_sram_req && !data_sram_addr_ok) || // data_sram bus request
    (data_sram_addr_ok_r && lsu_stall_wb);  // wb_stall
  assign s0_out_valid  = lsu_valid_i && (lsu_write_mem || lsu_read_mem) && !lsu_ex_taken;
  assign s0_out_ld_op  = lsu_ld_op;
  assign s0_out_offset = addr[1:0];
  assign s0_out_rt     = lsu_rt;
  assign s0_out        = {
    s0_out_ld_op,
    s0_out_offset,
    s0_out_rt
  };
  
  always@(posedge clk) begin
    if (rst) begin
      data_sram_addr_ok_r <= 1'b0;
    end else if (data_sram_addr_ok && lsu_stall_wb) begin
      // flipped up only when request hand shaked but wb is stalled
      data_sram_addr_ok_r <= 1'b1;
    end else if (data_sram_addr_ok_r && !lsu_stall_wb) begin
      data_sram_addr_ok_r <= 1'b0;
    end
  end

  assign data_sram_req      = 
    (data_vaddr_valid && data_paddr_valid) && !data_sram_addr_ok_r && (lsu_read_mem || lsu_write_mem) && !lsu_ex_taken;
  assign data_sram_uncache  = data_uncache;
  assign data_sram_wr       = lsu_write_mem;
  assign data_sram_size     = (lsu_ld_op[`OP_LBU] || lsu_ld_op[`OP_LB]) ? 2'b00 : 2'b10;
  assign data_sram_addr     = data_paddr;
  assign data_sram_wstrb[0] =
    lsu_st_op[`OP_SB ] && !offset[1] && !offset[0] ||
    lsu_st_op[`OP_SH ] && !offset[1] && !offset[0] ||
    lsu_st_op[`OP_SW ] ||
    lsu_st_op[`OP_SWL] ||
    lsu_st_op[`OP_SWR] && !offset[1] && !offset[0];
  assign data_sram_wstrb[1] =
    lsu_st_op[`OP_SB ] && !offset[1] && offset[0] ||
    lsu_st_op[`OP_SH ] && !offset[1] && !offset[0] ||
    lsu_st_op[`OP_SW ] ||
    lsu_st_op[`OP_SWL] && (offset[1] || offset[0]) ||
    lsu_st_op[`OP_SWR] && !offset[1];
  assign data_sram_wstrb[2] =
    lsu_st_op[`OP_SB ] && offset[1] && !offset[0] ||
    lsu_st_op[`OP_SH ] && offset[1] && !offset[0] ||
    lsu_st_op[`OP_SW ] ||
    lsu_st_op[`OP_SWL] && offset[1] ||
    lsu_st_op[`OP_SWR] && (!offset[1] || !offset[0]);
  assign data_sram_wstrb[3] =
    lsu_st_op[`OP_SB ] && offset[1] && offset[0] ||
    lsu_st_op[`OP_SH ] && offset[1] && !offset[0] ||
    lsu_st_op[`OP_SW ] ||
    lsu_st_op[`OP_SWL] && offset[1] && offset[0] ||
    lsu_st_op[`OP_SWR] ;
  assign lsu_st_result[`OP_SB ] = offset[1] ? 
                                 (offset[0] ? {lsu_rt[7:0], 24'd0      } : { 8'd0, lsu_rt[7:0], 16'd0}) :
                                 (offset[0] ? {16'd0, lsu_rt[7:0], 8'd0} : {24'd0, lsu_rt[7:0]       }) ;
  assign lsu_st_result[`OP_SH ] = offset[1] ? {lsu_rt[15:0], 16'd0} : {16'd0, lsu_rt[15:0]};
  assign lsu_st_result[`OP_SWL] = offset[1] ? 
                                 (offset[0] ? lsu_rt                 : { 8'd0, lsu_rt[31: 8]}) :
                                 (offset[0] ? {16'd0, lsu_rt[31:16]} : {24'd0, lsu_rt[31:24]}) ;
  assign lsu_st_result[`OP_SW ] = lsu_rt;
  assign lsu_st_result[`OP_SWR] = offset[1] ? 
                                 (offset[0] ? {lsu_rt[ 7:0], 24'd0} : {lsu_rt[15:0], 16'd0}) :
                                 (offset[0] ? {lsu_rt[23:0],  8'd0} : lsu_rt               ) ;
  assign data_sram_wdata = {32{lsu_st_op[`OP_SB ]}} & lsu_st_result[`OP_SB ] |
                      {32{lsu_st_op[`OP_SH ]}} & lsu_st_result[`OP_SH ] |
                      {32{lsu_st_op[`OP_SWL]}} & lsu_st_result[`OP_SWL] |
                      {32{lsu_st_op[`OP_SW ]}} & lsu_st_result[`OP_SW ] |
                      {32{lsu_st_op[`OP_SWR]}} & lsu_st_result[`OP_SWR] ;
  
  assign lsu_adel = (
    lsu_ld_op[`OP_LH] && offset[0] ||
    lsu_ld_op[`OP_LHU] && offset[0] ||
    lsu_ld_op[`OP_LW] && (offset[0] || offset[1])
  ) && (lsu_read_mem || lsu_write_mem) && lsu_valid_i;
  assign lsu_ades = (
    lsu_st_op[`OP_SH] && offset[0] ||
    lsu_st_op[`OP_SW] && (offset[0] || offset[1])
  ) && (lsu_read_mem || lsu_write_mem) && lsu_valid_i;
  assign lsu_tlbrel   = data_tlb_refill && lsu_valid_i && lsu_read_mem;
  assign lsu_tlbres   = data_tlb_refill && lsu_valid_i && lsu_write_mem;
  assign lsu_tlbil    = data_tlb_invalid && lsu_valid_i && lsu_read_mem;
  assign lsu_tlbis    = data_tlb_invalid && lsu_valid_i && lsu_write_mem;
  assign lsu_tlbm     = data_tlb_modified && lsu_valid_i;
  assign lsu_badvaddr = addr;

  assign data_vaddr = {addr[31:2], (lsu_ld_op[`OP_LBU] || lsu_ld_op[`OP_LB]) ? addr[1:0] : 2'b00};
  assign is_store = lsu_write_mem;

  /* stage 1 */
  assign s1_flush     = lsu_stall_ex && !s1_stall;
  assign s1_stall     = s1_valid && !data_sram_data_ok;
  assign lsu_stall_wb = s1_stall;
  assign lsu_valid_o  = s1_valid && data_sram_data_ok;
  always@(posedge clk) begin
    if (rst || s1_flush) begin
      s1_valid <= 1'b0;
    end else if(!s1_stall) begin
      s1_valid <= s0_out_valid;
    end
  end
  always@(posedge clk) begin
    if (!s1_stall) begin
      s1_reg <= s0_out;
    end
  end

  assign {
    s1_ld_op,
    s1_offset,
    s1_rt
  } = s1_reg;

  assign lsu_ld_result[`OP_LB ] = s1_offset[1] ? 
                                 (s1_offset[0] ? {{24{data_sram_rdata[31]}}, data_sram_rdata[31:24]} : {{24{data_sram_rdata[23]}}, data_sram_rdata[23:16]}) :
                                 (s1_offset[0] ? {{24{data_sram_rdata[15]}}, data_sram_rdata[15: 8]} : {{24{data_sram_rdata[ 7]}}, data_sram_rdata[ 7: 0]}) ;
  assign lsu_ld_result[`OP_LH ] = s1_offset[1] ? {{16{data_sram_rdata[31]}}, data_sram_rdata[31:16]} : {{16{data_sram_rdata[15]}}, data_sram_rdata[15: 0]};
  assign lsu_ld_result[`OP_LW ] = data_sram_rdata;
  assign lsu_ld_result[`OP_LWL] = s1_offset[1] ? 
                                 (s1_offset[0] ? data_sram_rdata                        : {data_sram_rdata[23:0], s1_rt[ 7:0]}) :
                                 (s1_offset[0] ? {data_sram_rdata[15:0], s1_rt[15:0]} : {data_sram_rdata[ 7:0], s1_rt[23:0]}) ;
  assign lsu_ld_result[`OP_LBU] = s1_offset[1] ? 
                                 (s1_offset[0] ? {24'd0, data_sram_rdata[31:24]} : {24'd0, data_sram_rdata[23:16]}) :
                                 (s1_offset[0] ? {24'd0, data_sram_rdata[15: 8]} : {24'd0, data_sram_rdata[ 7: 0]}) ;
  assign lsu_ld_result[`OP_LHU] = s1_offset[1] ? {16'd0, data_sram_rdata[31:16]} : {16'd0, data_sram_rdata[15: 0]};
  assign lsu_ld_result[`OP_LWR] = s1_offset[1] ? 
                                 (s1_offset[0] ? {s1_rt[31: 8], data_sram_rdata[31:24]} : {s1_rt[31:16], data_sram_rdata[31:16]}) :
                                 (s1_offset[0] ? {s1_rt[31:24], data_sram_rdata[31: 8]} : data_sram_rdata                          ) ;
  assign lsu_ld_data = {32{s1_ld_op[`OP_LB ]}} & lsu_ld_result[`OP_LB ] |
                       {32{s1_ld_op[`OP_LH ]}} & lsu_ld_result[`OP_LH ] |
                       {32{s1_ld_op[`OP_LWL]}} & lsu_ld_result[`OP_LWL] |
                       {32{s1_ld_op[`OP_LW ]}} & lsu_ld_result[`OP_LW ] |
                       {32{s1_ld_op[`OP_LBU]}} & lsu_ld_result[`OP_LBU] |
                       {32{s1_ld_op[`OP_LHU]}} & lsu_ld_result[`OP_LHU] |
                       {32{s1_ld_op[`OP_LWR]}} & lsu_ld_result[`OP_LWR] ;

  assign data_vaddr_valid = lsu_valid_i && (lsu_read_mem || lsu_write_mem);

endmodule